import React from 'react';
import { Tree, TreeProps } from 'antd';
import { EventDataNode } from 'antd/lib/tree';

const { DirectoryTree } = Tree;

interface DataNode {
  title: string;
  key: string;
  isLeaf?: boolean;
  children?: DataNode[];
}

export function updateTreeData(
  list: DataNode[],
  key: string,
  children: DataNode[],
): DataNode[] {
  return list.map((node) => {
    if (node.key === key) {
      return {
        ...node,
        children,
      };
    }
    if (node.children) {
      return {
        ...node,
        children: updateTreeData(node.children, key, children),
      };
    }
    return node;
  });
}

interface IDataTree {
  treeData: DataNode[];
  loadData: (node: any) => any;
  onExpand?: (keys: string[]) => void;
  onSelect?: (keys: string[], info: any) => void;
}

const DataTree: React.FC<
  Omit<TreeProps, 'onSelect' | 'onExpand'> & IDataTree
> = (props) => {
  const { treeData = [], loadData, ...rest } = props;
  const onLoadData = (e: EventDataNode): Promise<void> => {
    return new Promise((resolve) => {
      loadData(e).finally(() => {
        resolve();
      });
    });
  };
  return <DirectoryTree loadData={onLoadData} treeData={treeData} {...rest} />;
};

export default DataTree;
